{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "da42da9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import random\n",
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "text = (\n",
    "       'Hello, how are you? I am Romeo.\\n'\n",
    "       'Hello, Romeo My name is Juliet. Nice to meet you.\\n'\n",
    "       'Nice meet you too. How are you today?\\n'\n",
    "       'Great. My baseball team won the competition.\\n'\n",
    "       'Oh Congratulations, Juliet\\n'\n",
    "       'Thanks you Romeo'\n",
    "   )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "509f98ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[26, 22, 13, 24, 9, 28, 20], [26, 20, 23, 11, 15, 7, 6, 5, 12, 24], [6, 12, 24, 8, 22, 13, 24, 16], [4, 23, 19, 21, 14, 17, 25], [18, 27, 7], [10, 24, 20]]\n"
     ]
    }
   ],
   "source": [
    "sentences = re.sub(\"[.,!?\\\\-]\", '', text.lower()).split('\\n')  # filter '.', ',', '?', '!'\n",
    "word_list = list(set(\" \".join(sentences).split()))\n",
    "\n",
    "word_dict = {'[PAD]': 0, '[CLS]': 1, '[SEP]': 2, '[MASK]': 3}\n",
    "for i, w in enumerate(word_list):\n",
    "    word_dict[w] = i + 4\n",
    "    number_dict = {i: w for i, w in enumerate(word_dict)}\n",
    "    vocab_size = len(word_dict)\n",
    "\n",
    "token_list = [] #句子列表\n",
    "for i,sent in enumerate(sentences):\n",
    "    tmp =[]\n",
    "    for token in sent.split(\" \"):\n",
    "        tmp.append(word_dict.get(token))\n",
    "    token_list.append(tmp)\n",
    "print(token_list)\n",
    "\n",
    "# print(sentences)\n",
    "# tokens_a_index, tokens_b_index = random.randrange(len(sentences)), random.randrange(len(sentences)) \n",
    "# # print(tokens_a_index)\n",
    "# # print(tokens_b_index)\n",
    "\n",
    "# tokens_a, tokens_b = sentences[tokens_a_index], sentences[tokens_b_index]\n",
    "# print(tokens_a,word_dict['[CLS]'])\n",
    "# input_ids = [word_dict['[CLS]']] + tokens_a + [word_dict['[SEP]']] + tokens_b + [word_dict['[SEP]']]\n",
    "# print(input_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2f3fd57",
   "metadata": {},
   "source": [
    "### BERT的输入：\n",
    "![BERT_input](images/BERT_input.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b7859335",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_ids: [1, 26, 3, 23, 11, 15, 7, 6, 3, 12, 24, 2, 18, 27, 7, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] 32\n",
      "segment_ids: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] 32\n",
      "masked_tokens: [5, 20, 0, 0, 0] 5\n",
      "masked_pos: [8, 2, 0, 0, 0] 5\n",
      "label: False\n"
     ]
    }
   ],
   "source": [
    "batch_size = 4 #批次大小\n",
    "max_pred = 5 #一行输入里，对多有几个token要被mask\n",
    "\n",
    "d_model = 128 # 词嵌入的维度(和位置嵌入，segment嵌入的维度一致) embedding_dim\n",
    "max_seq_len = 32 # 每样本输入的最大token数\n",
    "n_segments = 2 # 句子数\n",
    "d_k = d_v = 64 # dimension of K(=Q), V, 不需要跟词嵌入的维度一致,\n",
    "# d_v = 68  # d_k,d_v必须相等？不想等意味着什么,dv>dk,dv<dk会发生什么？\n",
    "n_heads = 8 # number of heads in Multi-Head Attention\n",
    "n_layers = 6 # number of Encoder Layer\n",
    "d_ff = 2048 # FeedForward dimension\n",
    "\n",
    "def make_batch():\n",
    "    \"\"\"生成bert的训练数据\"\"\"\n",
    "    batch = []\n",
    "    positive = negative = 0\n",
    "    while positive != batch_size/2 or negative != batch_size/2:\n",
    "        tokens_a_index, tokens_b_index = random.randrange(len(sentences)), random.randrange(len(sentences)) \n",
    "        \n",
    "        #token_list,句子集合\n",
    "        tokens_a, tokens_b = token_list[tokens_a_index], token_list[tokens_b_index]\n",
    "        \n",
    "        #1.输入的token ids\n",
    "        input_ids = [word_dict['[CLS]']] + tokens_a + [word_dict['[SEP]']] + tokens_b + [word_dict['[SEP]']]\n",
    "\n",
    "        #2.句子的id，上一句，每个token对应位置都是0，下半句都是1\n",
    "        segment_ids = [0] * (1 + len(tokens_a) + 1) + [1] * (len(tokens_b) + 1)\n",
    "\n",
    "        #3. MASK LM \n",
    "        # 一个输入里面：15%的wordpiece token会被随机掩盖，\n",
    "        # 这15%的token中80%用[MASK]这个token来代替，10%用随机的一个词来替换,剩下的10%保持不变\n",
    "        mask_cnt = max(1, int(round(len(input_ids) * 0.15))) # 15 % of tokens in one sentence\n",
    "#         print(\"mask_cnt,input_ids_len:\",mask_cnt,len(input_ids))\n",
    "        n_pred =  min(max_pred, mask_cnt) \n",
    "        # 能够被用来mask的位置集合\n",
    "        cand_maked_pos = [i for i, token in enumerate(input_ids)\n",
    "                         if token != word_dict['[CLS]'] and token != word_dict['[SEP]']]\n",
    "#         print(\"cand_maked_pos:\",cand_maked_pos)\n",
    "        random.shuffle(cand_maked_pos)\n",
    "#         print(\"cand_maked_pos shuf:\",cand_maked_pos)\n",
    "        masked_tokens, masked_pos = [], []\n",
    "        for pos in cand_maked_pos[:n_pred]:\n",
    "            masked_pos.append(pos)\n",
    "            masked_tokens.append(input_ids[pos])\n",
    "            if random.random() < 0.8:  # 80%用[MASK]这个token来代替\n",
    "                input_ids[pos] = word_dict['[MASK]'] # make mask\n",
    "            elif random.random() < 0.5:  # 10%用随机的一个词来替换\n",
    "                index = random.randint(0, vocab_size - 1) # random index in vocabulary\n",
    "                input_ids[pos] = word_dict[number_dict[index]] # replace\n",
    "\n",
    "        # Zero Paddings, 其余部分补0\n",
    "        n_pad = max_seq_len - len(input_ids)\n",
    "        input_ids.extend([0] * n_pad)\n",
    "        segment_ids.extend([0] * n_pad)\n",
    "\n",
    "        # Zero Padding (100% - 15%) tokens\n",
    "        if max_pred > n_pred:\n",
    "            n_pad = max_pred - n_pred\n",
    "            masked_tokens.extend([0] * n_pad)\n",
    "            masked_pos.extend([0] * n_pad)\n",
    "\n",
    "        if tokens_a_index + 1 == tokens_b_index and positive < batch_size/2:\n",
    "            # 这种构造正负样本的方式简单粗暴，仅做演示用，太粗暴了，不过我喜欢\n",
    "            batch.append([input_ids, segment_ids, masked_tokens, masked_pos, True]) # IsNext\n",
    "            positive += 1\n",
    "        elif tokens_a_index + 1 != tokens_b_index and negative < batch_size/2:\n",
    "            batch.append([input_ids, segment_ids, masked_tokens, masked_pos, False]) # NotNext\n",
    "            negative += 1\n",
    "            \n",
    "    return batch\n",
    "\n",
    "\n",
    "batch = make_batch()\n",
    "# test\n",
    "print(\"input_ids:\",batch[0][0],len(batch[0][0]))\n",
    "print(\"segment_ids:\",batch[0][1],len(batch[0][1]))\n",
    "print(\"masked_tokens:\",batch[0][2],len(batch[0][2]))\n",
    "print(\"masked_pos:\",batch[0][3],len(batch[0][3]))\n",
    "print(\"label:\",batch[0][4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8db5009",
   "metadata": {},
   "source": [
    "![BERT模型](images/BERT-building-model.webp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4cae3239",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_attn_pad_mask(seq_q, seq_k):\n",
    "    \"\"\"\n",
    "    当输入小于max_len时，无效的区域如果为0，过softmax有值，\n",
    "    为了让这些无效的区域不参与运算，一般是给无效区域加一个很大的负数偏置\n",
    "    该函数目的，让padding部分变为true，其余为false，用于softmax前的负数填充 。#搜 masked_fill_\n",
    "    \"\"\"\n",
    "    batch_size, len_q = seq_q.size()  #len_q = input_ids's max_seq_len\n",
    "    batch_size, len_k = seq_k.size()\n",
    "    # eq(zero) is PAD token\n",
    "    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)  # [batch_size,1,len_k(=len_q)], one is masking\n",
    "    return pad_attn_mask.expand(batch_size, len_q, len_k)  # [batch_size,len_q,len_k]\n",
    "\n",
    "def gelu(x):\n",
    "    \"\"\"等价 torch.nn.functional.gelu\"\"\"\n",
    "    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))\n",
    "\n",
    "\n",
    "class Embedding(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Embedding, self).__init__()\n",
    "        self.tok_embed = nn.Embedding(vocab_size, d_model)  # token embedding\n",
    "        self.pos_embed = nn.Embedding(max_seq_len, d_model)  # position embedding\n",
    "        self.seg_embed = nn.Embedding(n_segments, d_model)  # segment embedding，上下句子\n",
    "        self.norm = nn.LayerNorm(d_model) #层归一化\n",
    "\n",
    "    def forward(self, x, seg):\n",
    "        seq_len = x.size(1) #x: token ids\n",
    "        pos = torch.arange(seq_len, dtype=torch.long) #0,1,...sel_len\n",
    "        pos = pos.unsqueeze(0).expand_as(x)  # (seq_len,) -> (batch_size, seq_len)\n",
    "        # embedding [batch_size,max_seq_len,d_model]\n",
    "        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)\n",
    "        return self.norm(embedding) \n",
    "\n",
    "    \n",
    "class ScaledDotProductAttention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ScaledDotProductAttention, self).__init__()\n",
    "\n",
    "    def forward(self, Q, K, V, attn_mask):\n",
    "        # scores : [batch_size,n_heads,len_q(=len_k),len_k(=len_q)]\n",
    "        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)\n",
    "        # 当输入小于max_len时，需要让这些无效的区域不参与运算，一般是给无效区域加一个很大的负数偏置\n",
    "        scores.masked_fill_(attn_mask, -1e9) # Fills elements of self tensor with value where mask is one.\n",
    "#         print(scores)\n",
    "        attn = nn.Softmax(dim=-1)(scores)\n",
    "#         print(attn)\n",
    "        context = torch.matmul(attn, V)\n",
    "#         print(\"shape:\",scores.shape,context.shape,attn.shape)\n",
    "        #[batch_size,head_cnt,len_q,d_k]\n",
    "        # scores:  torch.Size([4, 8, 32, 32]) \n",
    "        # attn:    torch.Size([4, 8, 32, 32])\n",
    "        # context: torch.Size([4, 8, 32, 68]) #d_k:32, d_v:68, 一般设定dk=dv，这里测试用 \n",
    "        # [batch_size,head_cnt,len_q,d_v]\n",
    "        return scores, context, attn\n",
    "    \n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.W_Q = nn.Linear(d_model, d_k * n_heads)\n",
    "        self.W_K = nn.Linear(d_model, d_k * n_heads)\n",
    "        self.W_V = nn.Linear(d_model, d_v * n_heads)\n",
    "        self.scaled_dot = ScaledDotProductAttention()\n",
    "        self.fc = nn.Linear(n_heads * d_v, d_model)\n",
    "        self.layernorm = nn.LayerNorm(d_model)\n",
    "\n",
    "    def forward(self, Q, K, V, attn_mask):\n",
    "        # q: [batch_size,seq_len,d_model], k: [batch_size,d_k,d_model], v: [batch_size,d_k,d_model]\n",
    "        residual, batch_size = Q, Q.size(0)\n",
    "        \n",
    "        # print(\"self.W_Q(Q).size():\",self.W_Q(Q).size())\n",
    "        # torch.Size([4, 32, 512]) d_k * n_heads:512\n",
    "        \n",
    "        # 需要将线性变换后多头部分单拆一维出来\n",
    "        # (batch_size, seq_len, d_model) -proj-> (B, S, D) \n",
    "        # -> split-> (B, S, H, W) -trans-> (B, Head_cnt, Seq_len, W)\n",
    "        \n",
    "        # q_s: [batch_size,n_heads,len_q,d_k]\n",
    "        q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1,2) \n",
    "        # k_s: [batch_size,n_heads,len_k,d_k]\n",
    "        k_s = self.W_K(K).view(batch_size, -1, n_heads, d_k).transpose(1,2) \n",
    "        # v_s: [batch_size,n_heads,len_k,d_v]\n",
    "        v_s = self.W_V(V).view(batch_size, -1, n_heads, d_v).transpose(1,2)  \n",
    "        \n",
    "        # attn_mask : [batch_size,n_heads,len_q,len_k]\n",
    "        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) \n",
    "        \n",
    "        # context: [batch_size,n_heads,len_q,d_v], attn: [batch_size,n_heads,len_q(=len_k),len_k(=len_q)]\n",
    "        scores, context, attn = self.scaled_dot(q_s, k_s, v_s, attn_mask)\n",
    "#         print(context.shape)\n",
    "        # context: [batch_size, len_q, n_heads*d_v]\n",
    "        context = context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v) \n",
    "        \n",
    "        output = self.fc(context) \n",
    "#         print(output)\n",
    "        # output: [batch_size,len_q,d_model]\n",
    "        return self.layernorm(output + residual), attn \n",
    "\n",
    "\n",
    "class PoswiseFeedForwardNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(PoswiseFeedForwardNet, self).__init__()\n",
    "        # 使用1*1卷积，将d_model维度转为d_ff维度\n",
    "        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1) \n",
    "        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)\n",
    "        self.layer_norm = nn.LayerNorm(d_model)\n",
    "        self.activ1 = nn.ReLU()\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        residual = inputs # inputs : [batch_size, len_q, d_model]\n",
    "        output = self.activ1(self.conv1(inputs.transpose(1, 2)))\n",
    "        output = self.conv2(output).transpose(1, 2)\n",
    "#         print(\"PoswiseFeedForwardNet.shape:\",output.shape)\n",
    "        return self.layer_norm(output + residual)\n",
    "\n",
    "\n",
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.enc_self_attn = MultiHeadAttention()\n",
    "        self.pos_ffn = PoswiseFeedForwardNet()\n",
    "\n",
    "    def forward(self, enc_inputs, enc_self_attn_mask):\n",
    "        # enc_inputs to same Q,K,V\n",
    "        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask) \n",
    "        # enc_outputs: [batch_size x len_q x d_model]\n",
    "        enc_outputs = self.pos_ffn(enc_outputs) \n",
    "        return enc_outputs, attn\n",
    "\n",
    "    \n",
    "class BERT(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BERT, self).__init__()\n",
    "        self.embedding = Embedding()\n",
    "        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])\n",
    "        self.fc = nn.Linear(d_model, d_model) #？\n",
    "        self.activ1 = nn.Tanh()\n",
    "        self.linear = nn.Linear(d_model, d_model)\n",
    "        self.activ2 = gelu\n",
    "        self.norm = nn.LayerNorm(d_model)\n",
    "        self.classifier = nn.Linear(d_model, 2)\n",
    "        # decoder is shared with embedding layer\n",
    "        embed_weight = self.embedding.tok_embed.weight\n",
    "        n_vocab, n_dim = embed_weight.size()\n",
    "        self.decoder = nn.Linear(n_dim, n_vocab, bias=False)\n",
    "        self.decoder.weight = embed_weight\n",
    "        self.decoder_bias = nn.Parameter(torch.zeros(n_vocab))\n",
    "\n",
    "    def forward(self, input_ids, segment_ids, masked_pos):\n",
    "        # output : [batch_size, seq_len, d_model]\n",
    "        output = self.embedding(input_ids, segment_ids)\n",
    "        enc_self_attn_mask = get_attn_pad_mask(input_ids, input_ids)\n",
    "        for layer in self.layers:\n",
    "            output, enc_self_attn = layer(output, enc_self_attn_mask)\n",
    "        # output : [batch_size, len, d_model], attn : [batch_size, n_heads, d_mode, d_model]\n",
    "        # it will be decided by first token(CLS)\n",
    "        h_pooled = self.activ1(self.fc(output[:, 0])) # [batch_size, d_model]\n",
    "        \n",
    "        #预测是否为上下句，输出维度：[batch_size, 2]\n",
    "        logits_clsf = self.classifier(h_pooled) \n",
    "        #masked_pos:[batch_size,max_pred,d_model]\n",
    "        masked_pos = masked_pos[:, :, None].expand(-1, -1, output.size(-1)) # [batch_size, max_pred, d_model]\n",
    "\n",
    "        # get masked position from final output of transformer.\n",
    "        h_masked = torch.gather(output, 1, masked_pos) # masking position [batch_size, max_pred, d_model]\n",
    "        h_masked = self.norm(self.activ2(self.linear(h_masked)))\n",
    "        logits_lm = self.decoder(h_masked) + self.decoder_bias # [batch_size, max_pred, n_vocab]\n",
    "\n",
    "        return logits_lm, logits_clsf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "84fbba24",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "self.W_Q(Q).size(): torch.Size([4, 32, 512])\n",
      "self.W_Q(Q).size(): torch.Size([4, 32, 512])\n",
      "self.W_Q(Q).size(): torch.Size([4, 32, 512])\n",
      "self.W_Q(Q).size(): torch.Size([4, 32, 512])\n",
      "self.W_Q(Q).size(): torch.Size([4, 32, 512])\n",
      "self.W_Q(Q).size(): torch.Size([4, 32, 512])\n",
      "self.W_Q(Q).size(): torch.Size([4, 32, 512])\n",
      "self.W_Q(Q).size(): torch.Size([4, 32, 512])\n"
     ]
    }
   ],
   "source": [
    "# 测试模型里的各个组件\n",
    "input_ids, segment_ids, masked_tokens, masked_pos, isNext = map(torch.LongTensor, zip(*batch))\n",
    "\n",
    "seq_len = input_ids.size(1)\n",
    "# print(x.shape)\n",
    "pos = torch.arange(seq_len, dtype=torch.long)\n",
    "pos = pos.unsqueeze(0).expand_as(input_ids) \n",
    "# print(pos.shape)\n",
    "\n",
    "pad_attn_mask = input_ids.data.eq(0).unsqueeze(1)\n",
    "# print(pad_attn_mask)\n",
    "\n",
    "enc_self_attn_mask = get_attn_pad_mask(input_ids, input_ids)\n",
    "# print(enc_self_attn_mask)\n",
    "# print(enc_self_attn_mask.shape) #[batch_size,max_seq_len,max_seq_len]\n",
    "attn_mask = enc_self_attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)  #[batch_size,n_heads,max_seq_len,max_seq_len]\n",
    "# print(attn_mask.shape)\n",
    "# print(attn_mask)\n",
    "\n",
    "embedding = Embedding()\n",
    "enc_inputs = embedding(input_ids, segment_ids)\n",
    "# # print(enc_inputs)\n",
    "\n",
    "mha = MultiHeadAttention()\n",
    "ret = mha(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask)\n",
    "# # print(ret)\n",
    "\n",
    "el = EncoderLayer()\n",
    "ret = el(enc_inputs,enc_self_attn_mask)\n",
    "# # print(ret)\n",
    "\n",
    "model = BERT()\n",
    "logits_lm, logits_clsf = model(input_ids, segment_ids, masked_pos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "da5fe2db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0010 cost = 5.890497\n",
      "Epoch: 0020 cost = 3.496031\n",
      "Epoch: 0030 cost = 2.573384\n",
      "Epoch: 0040 cost = 2.201141\n",
      "Epoch: 0050 cost = 2.094465\n",
      "Epoch: 0060 cost = 2.077641\n",
      "Epoch: 0070 cost = 2.063523\n",
      "Epoch: 0080 cost = 2.062627\n",
      "Epoch: 0090 cost = 2.059888\n",
      "Epoch: 0100 cost = 2.059503\n",
      "Hello, how are you? I am Romeo.\n",
      "Hello, Romeo My name is Juliet. Nice to meet you.\n",
      "Nice meet you too. How are you today?\n",
      "Great. My baseball team won the competition.\n",
      "Oh Congratulations, Juliet\n",
      "Thanks you Romeo\n",
      "['[CLS]', 'hello', '[MASK]', 'my', 'name', 'is', 'juliet', 'nice', '[MASK]', 'meet', 'you', '[SEP]', 'oh', 'congratulations', 'juliet', '[SEP]']\n",
      "masked tokens list :  [5, 20]\n",
      "predict masked tokens list :  []\n",
      "isNext :  False\n",
      "predict isNext :  False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[W NNPACK.cpp:79] Could not initialize NNPACK! Reason: Unsupported hardware.\n"
     ]
    }
   ],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "for epoch in range(100):\n",
    "    optimizer.zero_grad()\n",
    "    logits_lm, logits_clsf = model(input_ids, segment_ids, masked_pos)\n",
    "    loss_lm = criterion(logits_lm.transpose(1, 2), masked_tokens) # for masked LM\n",
    "    loss_lm = (loss_lm.float()).mean()\n",
    "    loss_clsf = criterion(logits_clsf, isNext) # for sentence classification\n",
    "    loss = loss_lm + loss_clsf\n",
    "    if (epoch + 1) % 10 == 0:\n",
    "        print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "# Predict mask tokens\n",
    "input_ids, segment_ids, masked_tokens, masked_pos, isNext = map(torch.LongTensor, zip(batch[0]))\n",
    "print(text)\n",
    "print([number_dict[w.item()] for w in input_ids[0] if number_dict[w.item()] != '[PAD]'])\n",
    "\n",
    "logits_lm, logits_clsf = model(input_ids, segment_ids, masked_pos)\n",
    "logits_lm = logits_lm.data.max(2)[1][0].data.numpy()\n",
    "print('masked tokens list : ',[pos.item() for pos in masked_tokens[0] if pos.item() != 0])\n",
    "print('predict masked tokens list : ',[pos for pos in logits_lm if pos != 0])\n",
    "\n",
    "logits_clsf = logits_clsf.data.max(1)[1].data.numpy()[0]\n",
    "print('isNext : ', True if isNext else False)\n",
    "print('predict isNext : ',True if logits_clsf else False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85190ee5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
